// Copyright (c) 2024 by Rockchip Electronics Co., Ltd. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>

#include <sched.h>
#include <sys/syscall.h>

#include <time.h>
#include <sys/time.h>
#include "yolo11.h"
#include "common.h"
#include "file_utils.h"
#include "image_utils.h"

using namespace std;
float nms_threshold = 0.45; // 默认的NMS阈值
float box_conf_threshold = 0.25; // 默认的置信度阈值
float resizeBate = 1.0f;
static void dump_tensor_attr(rknn_tensor_attr *attr)
{
    printf("  index=%d, name=%s, n_dims=%d, dims=[%d, %d, %d, %d], n_elems=%d, size=%d, fmt=%s, type=%s, qnt_type=%s, "
           "zp=%d, scale=%f\n",
           attr->index, attr->name, attr->n_dims, attr->dims[0], attr->dims[1], attr->dims[2], attr->dims[3],
           attr->n_elems, attr->size, get_format_string(attr->fmt), get_type_string(attr->type),
           get_qnt_type_string(attr->qnt_type), attr->zp, attr->scale);
}

pid_t gettid() {
	return syscall(SYS_gettid);
}

void cpu_num_set(unsigned char num)
{
	// pid_t pid = gettid(); // LWP
	// pthread_t tid = pthread_self();

	// long cpu_num = sysconf(_SC_NPROCESSORS_ONLN);
	// if (cpu_num <= 0) {
	// 	printf("pid=%d,tid=%lu,get cpu info error:%ld\n",
	// 			pid, tid, cpu_num);
	// 	return;
	// }

	cpu_set_t cpu_set;
	CPU_ZERO(&cpu_set);
	CPU_SET(num, &cpu_set);

	// int res = pthread_setaffinity_np(tid, sizeof(cpu_set), &cpu_set);
    int res = sched_setaffinity(0, sizeof(cpu_set), &cpu_set);
	if (res != 0) {
		printf("tid=0,pthread_setaffinity_np error:%d(%s)\n",res, strerror(res));
		return;
	}

	// printf("pid=%d,tid=%lu,bind cpu[%d] ok!\n",
	// 		pid, tid, num);
}

int testnum = 0;

int inference_yolo11_model(rknn_app_context_t *app_ctx, image_buffer_t *img, object_detect_result_list *od_results) {
    int ret;
    image_buffer_t dst_img;
    letterbox_t letter_box;
    struct timeval start_time, end_time;
    struct timeval start_time1, end_time1;
    float timeuse = 0.0f;
    int bg_color = 0; // 114;
    gettimeofday(&start_time1, NULL);
    if ((!app_ctx) || !(img) || (!od_results)) {
        printf("gettimeofday(&start_time1, NULL) fail!\n");
        return -1;
    }

    rknn_input inputs[app_ctx->io_num.n_input];
    rknn_output outputs[app_ctx->io_num.n_output];

    memset(od_results, 0x00, sizeof(*od_results));
    memset(&letter_box, 0, sizeof(letterbox_t));
    memset(&dst_img, 0, sizeof(image_buffer_t));
    memset(inputs, 0, sizeof(inputs));
    memset(outputs, 0, sizeof(outputs));

    // Pre Process
    dst_img.width = app_ctx->model_width;
    dst_img.height = app_ctx->model_height;
    dst_img.format = IMAGE_FORMAT_RGB888;
    dst_img.size = get_image_size(&dst_img);
    dst_img.virt_addr = (unsigned char *)malloc(dst_img.size);

    if (dst_img.virt_addr == NULL && dst_img.fd == 0) {
        printf("malloc buffer size:%d fail!\n", dst_img.size);
        return -1;
    }
    // gettimeofday(&start_time, NULL);
    // letterbox
    ret = convert_image_with_letterbox(img, &dst_img, &letter_box, bg_color);
    if (ret < 0) {
        printf("convert_image_with_letterbox fail! ret=%d\n", ret);
        return -1;
    }
    static int testimg = 0;
    /*
    if (testimg == 20)
    {
        FILE* fptest = fopen("./testimg.bin", "wb");
        fwrite(dst_img.virt_addr, 1, app_ctx->model_width * app_ctx->model_height * app_ctx->model_channel, fptest);
        fclose(fptest);
        sync();
    }
    testimg++;
    */
    // gettimeofday(&end_time, NULL);
    // timeuse = (end_time.tv_sec - start_time.tv_sec) * 1000 + ((float)end_time.tv_usec - start_time.tv_usec) / 1000; //ms
    // printf("get convert img time %0.2f\n", timeuse);

    // gettimeofday(&start_time, NULL);
    // Set Input Data
    inputs[0].index = 0;
    inputs[0].type = RKNN_TENSOR_UINT8;
    inputs[0].fmt = RKNN_TENSOR_NHWC;
    inputs[0].size = app_ctx->model_width * app_ctx->model_height * app_ctx->model_channel;
    inputs[0].buf = dst_img.virt_addr;
    ret = rknn_inputs_set(app_ctx->rknn_ctx, app_ctx->io_num.n_input, inputs);
    if (ret < 0)
    {
        printf("rknn_input_set fail! ret=%d\n", ret);
        return -1;
    }
    // gettimeofday(&end_time, NULL);
    // timeuse = (end_time.tv_sec - start_time.tv_sec) * 1000 + ((float)end_time.tv_usec - start_time.tv_usec) / 1000; //ms
    // printf("input time %0.2f\n", timeuse);

    // gettimeofday(&start_time, NULL);
    // Run
    // printf("rknn_run\n");
    ret = rknn_run(app_ctx->rknn_ctx, nullptr);
    if (ret < 0) {
        printf("rknn_run fail! ret=%d\n", ret);
        return -1;
    }

    // gettimeofday(&end_time, NULL);
    // timeuse = (end_time.tv_sec - start_time.tv_sec) * 1000 + ((float)end_time.tv_usec - start_time.tv_usec) / 1000; //ms
    // printf("yolov11 time %0.2f\n", timeuse);

    // gettimeofday(&start_time, NULL);
     // Get Output
    memset(outputs, 0, sizeof(outputs));
    for (int i = 0; i < app_ctx->io_num.n_output; i++)
    {
        outputs[i].index = i;
        outputs[i].want_float = (!app_ctx->is_quant);
    }
    ret = rknn_outputs_get(app_ctx->rknn_ctx, app_ctx->io_num.n_output, outputs, NULL);
    if (ret < 0)
    {
        printf("rknn_outputs_get fail! ret=%d\n", ret);
        if (dst_img.virt_addr != NULL)
        {
            free(dst_img.virt_addr);
        }
        return -1;
    }
    //  gettimeofday(&end_time, NULL);
    // timeuse = (end_time.tv_sec - start_time.tv_sec) * 1000 + ((float)end_time.tv_usec - start_time.tv_usec) / 1000; //ms
    // printf("get NCHW time %0.2f\n", timeuse);

    // gettimeofday(&start_time, NULL);
    // Post Process
    post_process(app_ctx, outputs, &letter_box, box_conf_threshold, nms_threshold, od_results);
    //  gettimeofday(&end_time, NULL);
    // timeuse = (end_time.tv_sec - start_time.tv_sec) * 1000 + ((float)end_time.tv_usec - start_time.tv_usec) / 1000; //ms
    // printf("get post_process time %0.2f\n", timeuse);

    // Remeber to release rknn output
    rknn_outputs_release(app_ctx->rknn_ctx, app_ctx->io_num.n_output, outputs);
    if (dst_img.virt_addr != NULL)
    {
        free(dst_img.virt_addr);
    }

    gettimeofday(&end_time1, NULL);
    timeuse = (end_time1.tv_sec - start_time1.tv_sec) * 1000 + ((float)end_time1.tv_usec - start_time1.tv_usec) / 1000; //ms
    // printf("just yolo time %0.2f\n", timeuse);
    return ret;
}

namespace sv2
{
bool rknnYolov11::setup(std::string model_path, float confidence, float nms_thresh,  nlohmann::json g_dataset_categories,  nlohmann::json dataset_name)
{
    int ret;
    int model_len = 0;
    char *model;
    rknn_context ctx = 0;
    nms_threshold = nms_thresh;
    box_conf_threshold = confidence;
    this->_g_dataset_categories = g_dataset_categories;
    this->_dataset_name = dataset_name;
    app_ctx = new rknn_app_context_t;
    memset(app_ctx, 0, sizeof(rknn_app_context_t));
    printf("class num: %lu\n", this->_g_dataset_categories[this->_dataset_name].size());
    init_post_process(this->_g_dataset_categories[this->_dataset_name].size());
    cout << this->_g_dataset_categories[this->_dataset_name] << endl;

    // Load RKNN Model
    model_len = read_data_from_file(model_path.c_str(), &model);
    if (model == NULL) {
        printf("load_model fail!\n");
        return false;
    }
    ret = rknn_init(&ctx, model, model_len, 0, NULL);
    free(model);
    if (ret < 0) {
        printf("rknn_init fail! ret=%d\n", ret);
        return false;
    }
    // Get Model Input Output Number
    rknn_input_output_num io_num;
    ret = rknn_query(ctx, RKNN_QUERY_IN_OUT_NUM, &io_num, sizeof(io_num));
    if (ret != RKNN_SUCC) {
        printf("rknn_query fail! ret=%d\n", ret);
        return false;
    }
    printf("model input num: %d, output num: %d\n", io_num.n_input, io_num.n_output);

    // Get Model Input Info
    printf("input tensors:\n");
    rknn_tensor_attr input_attrs[io_num.n_input];
    memset(input_attrs, 0, sizeof(input_attrs));
    for (int i = 0; i < io_num.n_input; i++) {
        input_attrs[i].index = i;
        ret = rknn_query(ctx, RKNN_QUERY_INPUT_ATTR, &(input_attrs[i]), sizeof(rknn_tensor_attr));
        if (ret != RKNN_SUCC) {
            printf("rknn_query fail! ret=%d\n", ret);
            return false;
        }
        dump_tensor_attr(&(input_attrs[i]));
    }

    // Get Model Output Info
    printf("output tensors:\n");
    rknn_tensor_attr output_attrs[io_num.n_output];
    memset(output_attrs, 0, sizeof(output_attrs));
    for (int i = 0; i < io_num.n_output; i++) {
        output_attrs[i].index = i;
        ret = rknn_query(ctx, RKNN_QUERY_OUTPUT_ATTR, &(output_attrs[i]), sizeof(rknn_tensor_attr));
        if (ret != RKNN_SUCC) {
            printf("rknn_query fail! ret=%d\n", ret);
            return false;
        }
        dump_tensor_attr(&(output_attrs[i]));
    }

    // Set to context
    app_ctx->rknn_ctx = ctx;

    // TODO
    if (output_attrs[0].qnt_type == RKNN_TENSOR_QNT_AFFINE_ASYMMETRIC && output_attrs[0].type == RKNN_TENSOR_INT8) {
        app_ctx->is_quant = true;
    } else {
        app_ctx->is_quant = false;
    }
    printf("is_quant is %d\n", app_ctx->is_quant);

    app_ctx->io_num = io_num;
    app_ctx->input_attrs = (rknn_tensor_attr *)malloc(io_num.n_input * sizeof(rknn_tensor_attr));
    memcpy(app_ctx->input_attrs, input_attrs, io_num.n_input * sizeof(rknn_tensor_attr));
    app_ctx->output_attrs = (rknn_tensor_attr *)malloc(io_num.n_output * sizeof(rknn_tensor_attr));
    memcpy(app_ctx->output_attrs, output_attrs, io_num.n_output * sizeof(rknn_tensor_attr));

    if (input_attrs[0].fmt == RKNN_TENSOR_NCHW) {
        printf("model is NCHW input fmt\n");
        app_ctx->model_channel = input_attrs[0].dims[1];
        app_ctx->model_height = input_attrs[0].dims[2];
        app_ctx->model_width = input_attrs[0].dims[3];
    } else {
        printf("model is NHWC input fmt\n");
        app_ctx->model_height = input_attrs[0].dims[1];
        app_ctx->model_width = input_attrs[0].dims[2];
        app_ctx->model_channel = input_attrs[0].dims[3];
    }
    printf("model input height=%d, width=%d, channel=%d\n",
           app_ctx->model_height, app_ctx->model_width, app_ctx->model_channel);

    return true;
}
bool rknnYolov11::isOpened()
{
    if(app_ctx->model_height <= 0 || app_ctx->model_width <= 0 || app_ctx->model_channel <= 0)
    {
        return false;
    }
    return true;
}

void rknnYolov11::getconfig(sms::Publisher* _result_writer, sms::Publisher* _show_writer, sms::Publisher* _next_writer, bool b_use_shm, bool launch_next_emit)
{
    this->_result_writer = _result_writer;
    this->_show_writer = _show_writer;
    this->_next_writer = _next_writer;
    this->b_use_shm = b_use_shm;
    this->launch_next_emit = launch_next_emit;
    return;
}
void rknnYolov11::pushmsg(nlohmann::json msg, nlohmann::json msg_results, cv::Mat image)
{
    this->_result_writer->publish(msg_results);

    if(msg.contains("img_total") && this->launch_next_emit)
    {
        nlohmann::json next_msg = sms::def_msg("std_msgs::Boolean");
        next_msg["data"] = true;
        this->_next_writer->publish(next_msg);
    }
    if (this->b_use_shm)
    {
        msg = this->_show_writer->cvimg2sms_mem(image);
    }
    msg["spirecv_msgs::2DTargets"] = msg_results;

    this->_show_writer->publish(msg);
}
void rknnYolov11::trans_det_results(object_detect_result_list results, int w, int h, std::vector<int> roi, nlohmann::json& sms_results)
{
    nlohmann::json ann;
    nlohmann::json dataset_categories = this->_g_dataset_categories[this->_dataset_name];
    sms_results = sms::def_msg("spirecv_msgs::2DTargets");
    sms_results["file_name"] = "";
    sms_results["height"] = h;
    sms_results["width"] = w;
    sms_results["targets"] = nlohmann::json::array();
    float x1, y1, x2, y2;
    if(!roi.empty())
    {
        nlohmann::json roibak = {roi[0], roi[1], roi[2] - roi[0], roi[3] - roi[1]};
        sms_results["rois"] = {roibak};
    }
    for(int i = 0; i < results.count; i++)
    {
        if(!roi.empty())
        {
            results.results[i].box.left /= resizeBate;
            results.results[i].box.right /= resizeBate;
        }
        x1 = round(results.results[i].box.left * 1000) / 1000;
        y1 = round(results.results[i].box.top * 1000) / 1000;
        x2 = round(results.results[i].box.right * 1000) / 1000;
        y2 = round(results.results[i].box.bottom * 1000) / 1000;
        ann["category_name"] = dataset_categories[results.results[i].cls_id];//coco_cls_to_name(results.results[i].cls_id);
        ann["category_id"] = results.results[i].cls_id;
        ann["score"] = round(results.results[i].prop * 1000) / 1000;
        ann["bbox"] = {x1, y1, x2, y2};
        ann["bbox"][2] = x2-x1;
        ann["bbox"][3] = y2-y1;
        if(!roi.empty())
        {    
            ann["bbox"][0] = x1 + roi[0];
            ann["bbox"][1] = y1 + roi[1];
        }
         sms_results["targets"].push_back(ann);
    }
}


//!用于每一帧图像，在循环中被调用
/*!
\param image:
*/
void rknnYolov11::stream(cv::Mat image, nlohmann::json msg)
{
    // cpu_num_set(7);
    int ret = 0;
    vector<int> roi;
    cv::Mat imageuse;
    cv::Mat img_infer;
    cv::Mat img_test;
    nlohmann::json msg_results;

    // image_buffer_t src_image;
    // read_image("./dog_bike_car.jpg", &src_image); //guo test no
    // imageuse = cv::imread("/home/orangepi/spirecv-pro/spirecv/algorithm/common_det/Rknn/YOLOv11DetCppNode_Rknn/build/dog_bike_car.bmp", 1);
    // printf("image: %d %d %d\n", imageuse.cols, imageuse.rows, imageuse.channels());
    // if(imageuse.empty())
    // {
    //     printf("image error\n");
    //     return ;
    // }
    // printf("read a img\n");
    // cv::cvtColor(image, imageuse, cv::COLOR_BGR2RGB);
    if(msg.contains("rois") && ! msg["rois"].empty())
    {
        roi = msg["rois"][0].get<vector<int>>();
        cv::Rect crop_rect(roi[0], roi[1], roi[2] - roi[0], roi[3] - roi[1]);
#if 0
        img_infer = image(crop_rect).clone();
        printf("roi: %d %d %d %d\n", roi[0], roi[1], roi[2], roi[3]);
        if(img_infer.cols % 16 != 0)
        {
            resize(img_infer, img_test, cv::Size(img_infer.cols/16 * 16 + 16,  img_infer.rows));
            resizeBate = (float)img_test.cols / img_infer.cols;
            printf("src %dx%d dst %dx%d b %0.2f\n", img_infer.cols, img_infer.rows, img_test.cols, img_test.rows, resizeBate);
            img_infer = img_test;
        }
#else
        if(crop_rect.width %16 != 0)
        {
            if(crop_rect.x == 0)
            {
                crop_rect.width = crop_rect.width/16*16 + 16;
                if(crop_rect.width > image.cols)
                {
                    crop_rect.width = image.cols;
                }
            }
            else
            {
                crop_rect.width = crop_rect.width/16*16 + 16;
                if((crop_rect.width + crop_rect.x) > image.cols)
                {
                    crop_rect.x -= crop_rect.width + crop_rect.x  - image.cols;
                    roi[0] -= crop_rect.x;
                }
            }
            img_infer = image(crop_rect).clone();
        }
        else
        {
            img_infer = image(crop_rect).clone();
        }
#endif
    }
    else
    {
        img_infer = image;
    }
    string file_name = msg.value("file_name", "");
    image_buffer_t img;
    // unsigned char* outbuf = (unsigned char*)malloc(img_infer.cols * img_infer.rows * 3);
    // memcpy(outbuf, img_infer.data, img_infer.cols * img_infer.rows * 3);
    img.width = img_infer.cols;
    img.height = img_infer.rows;
    img.format = IMAGE_FORMAT_BGR888;
    img.size = img_infer.cols * img_infer.rows * 3;
    img.virt_addr = img_infer.data;
    // printf("img %dx%d, size %d format %d ,%p,  src%dx%d\n", img.width, img.height, img.size,img.format, img.virt_addr, image.cols, image.rows);
    object_detect_result_list od_results;
    ret = inference_yolo11_model(app_ctx, &img, &od_results); //guo test img
    if (ret != 0)
    {
        printf("inference_yolo11_model fail! ret=%d\n", ret);
        return;
    }

    trans_det_results(od_results, image.cols, image.rows, roi, msg_results);
    msg_results["file_name"] = file_name;
    msg_results["dataset"] = this->_dataset_name;
    if(msg.contains("client_id"))
    {
        msg_results["client_id"] = msg["client_id"];
    }
    if(msg.contains("img_id"))
    {
        msg_results["img_id"] = msg["img_id"];
    }
    if(msg.contains("img_total"))
    {
        msg_results["img_total"] = msg["img_total"];
    }
    
    pushmsg(msg, msg_results, image);

    // nlohmann::json dataset_categories = this->_g_dataset_categories[this->_dataset_name];
    // 画框和概率
    // char text[256];
    // for (int i = 0; i < od_results.count; i++)
    // {
    //     object_detect_result *det_result = &(od_results.results[i]);
    //     printf("%s @ (%d %d %d %d) %.3f\n", dataset_categories[det_result->cls_id].dump().c_str(),
    //            det_result->box.left, det_result->box.top,
    //            det_result->box.right, det_result->box.bottom,
    //            det_result->prop);
    //     int x1 = det_result->box.left;
    //     int y1 = det_result->box.top;
    //     int x2 = det_result->box.right;
    //     int y2 = det_result->box.bottom;

    //     cv::rectangle(image, cv::Point(x1, y1), cv::Point(x2, y2), cv::Scalar(0, 255, 255));
    //     // draw_rectangle(&src_image, x1, y1, x2 - x1, y2 - y1, COLOR_BLUE, 3);

    //     // sprintf(text, "%s %.1f%%", coco_cls_to_name(det_result->cls_id), det_result->prop * 100);
    //     // draw_text(&src_image, text, x1, y1 - 20, COLOR_RED, 10);
    // }
    // free(outbuf);
}
//!释放资源
void rknnYolov11::release()
{
    int ret;
    deinit_post_process();
    if (app_ctx->input_attrs != NULL) {
        free(app_ctx->input_attrs);
        app_ctx->input_attrs = NULL;
    }
    if (app_ctx->output_attrs != NULL) {
        free(app_ctx->output_attrs);
        app_ctx->output_attrs = NULL;
    }
    if (app_ctx->rknn_ctx != 0) {
        ret = rknn_destroy(app_ctx->rknn_ctx);
        if (ret != RKNN_SUCC) {
            printf("rknn_destroy fail! ret=%d\n", ret);
            return;
        }
        app_ctx->rknn_ctx = 0;

    }
    return;
}
}